#lang racket

;LEAF OPEATIONS-----------------------------------------------------------------------------------------------

(define (make-leaf symbol weight)
  (list 'leaf symbol weight))

(define (leaf? object)
  (eq? (car object) 'leaf))

(define (symbol-leaf x) (cadr x))

(define (weight-leaf x) (caddr x))

;TREE OPERATIONS----------------------------------------------------------------------------------------------

(define (make-code-tree left right)
  (list left
        right
        (append (symbols left) (symbols right))
        (+ (weight left) (weight right))))

(define (left-branch tree) (car tree))

(define (right-branch tree) (cadr tree))

(define (symbols tree)
  (if (leaf? tree)
      (list (symbol-leaf tree))
      (caddr tree)))

(define (weight tree)
  (if (leaf? tree)
      (weight-leaf tree)
      (cadddr tree)))

(define (adjoin-set x set)
  (cond ((null? set) (list x))
        ((< (weight x) (weight (car set))) (cons x set))
        (else (cons (car set)
                    (adjoin-set x (cdr set))))))

(define (make-leaf-set pairs)
  (if (null? pairs)
      '()
      (let ((pair (car pairs)))
        (adjoin-set (make-leaf (car pair)
                               (cadr pair))
                    (make-leaf-set (cdr pairs))))))

(define (merge-nodes pos nodes)
  (cond ((or (null? nodes) (null? (cdr nodes)) (<= pos 0)) (error "unexpected end of list: "))
        ((= pos 1) (cons (make-code-tree (car nodes) (cadr nodes)) (cddr nodes)))
        (else (cons (car nodes) (merge-nodes (- pos 1) (cdr nodes))))))

(define (find-least-pair-of-nodes nodes act-min min-num act-num)
  (cond ((or (null? nodes)
             (null? (cdr nodes))) min-num)
        (else (let ((p1 (car nodes))
                    (p2 (cadr nodes)))
                (let ((sum (+ (weight p1) (weight p2))))
                  (if (< sum act-min)
                      (find-least-pair-of-nodes (cdr nodes) sum act-num (+ act-num 1))
                      (find-least-pair-of-nodes (cdr nodes) act-min min-num (+ act-num 1))))))))

(define (find-smallest-pair nodes)
  (find-least-pair-of-nodes nodes -1 1 1))

(define (successive-merge nodes)
  (define (merge-one-pair pairs)
    (merge-nodes (find-smallest-pair pairs) pairs))
  (if (or (null? nodes) (null? (cdr nodes)))
      (car nodes)
      (successive-merge (merge-one-pair nodes))))

(define (generate-huffman-tree pairs)
  (successive-merge (make-leaf-set pairs)))

;DECODE OPERATIONS--------------------------------------------------------------------------------------------

(define (decode bits tree)
  (define (decode-1 bits current-branch)
    (if (null? bits)
        '()
        (let ((next-branch
               (choose-branch (car bits) current-branch)))
          (if (leaf? next-branch)
              (cons (symbol-leaf next-branch)
                    (decode-1 (cdr bits) tree))
              (decode-1 (cdr bits) next-branch)))))
  (decode-1 bits tree))

(define (choose-branch bit branch)
  (cond ((= bit 0) (left-branch branch))
        ((= bit 1) (right-branch branch))
        (else (error "bad bit: " bit))))

;ENCODE OPERATIONS--------------------------------------------------------------------------------------------

(define (encode message tree)
  (if (null? message)
      '()
      (append (encode-symbol (car message) tree)
              (encode (cdr message) tree))))

(define (encode-symbol symbol tree)
  (let ((result (find-leaf symbol tree '())))
    (if (not (eq? result false))
        (reverse result)
        (error "symbol not in tree: " symbol))))

(define (find-leaf symbol tree code)
  (if (leaf? tree) (if (eq? (symbol-leaf tree) symbol)
                       code
                       false)
      (find-branch symbol tree code)))

(define (find-branch symbol tree code)
  (let ((left-result (find-leaf symbol (left-branch tree) (cons 0 code))))
    (if (not (eq? left-result false))
        left-result
        (let ((right-result (find-leaf symbol (right-branch tree) (cons 1 code))))
          (if (not (eq? right-result false))
              right-result
              false)))))

;TESTING------------------------------------------------------------------------------------------------------

(define leaf-set '((A 2) (BOOM 1) (GET 2) (JOB 2)
                         (NA 16) (SHA 3) (YIP 9)
                         (WAH 1)))

(define tree (generate-huffman-tree leaf-set))

(define message '(GET A JOB
                      SHA NA NA NA NA NA NA NA NA
                      GET A JOB
                      SHA NA NA NA NA NA NA NA NA
                      WAH YIP YIP YIP YIP YIP YIP YIP YIP YIP
                      SHA BOOM))

(define encoded-message (encode message tree))

(length encoded-message)

(decode encoded-message tree)